import os

import numpy as np
import matplotlib
from matplotlib import pyplot as plt, ticker

from utils import load_variable
matplotlib.use('QtAgg')
# plt.rcParams['text.usetex'] = True


def save_lazyselect_analysis_plot(_result, _filepath, _filepath_time):
    fig, ax = plt.subplots()
    # ax2 = ax.twinx()  # instantiate a second axes that shares the same x-axis

    sampling_ratios = np.array(_result['sampling_ratios'])
    avg_time_consumes = np.array([each * 1000 for each in _result['avg_time_consumes']])
    success_rates = np.array(_result['success_rates'])
    fail_rates1 = np.array(_result['fail_rates']['k-th element is not in P'])
    fail_rates2 = np.array(_result['fail_rates']['|p| >= 4 * n ^ sampling_ratio'])
    # ax2.yaxis.tick_left()
    # ax2.yaxis.set_label_position("left")
    # ax.yaxis.tick_right()
    # ax.yaxis.set_label_position("right")
    width = 1 / 8 * 0.8
    ax.bar(sampling_ratios, success_rates, width=width, color='#2874A6', label='success')
    ax.bar(sampling_ratios, fail_rates1, width=width, color='#D2691E',
           label=r'$k\notin[Rank_s(a),Rank_s(b)]$', bottom=success_rates)
    ax.bar(sampling_ratios, fail_rates2, width=width, color='#2E8B57',
           label=r'$|P|\geq4n^r$', bottom=fail_rates1+success_rates)
    ax.set_ylabel('happen rate')
    # ax.tick_params(axis='y', labelcolor='#2874A6')
    positions = [i / 10 for i in range(11)]
    labels = [str(int(each * 100)) + '%' for each in positions]
    ax.yaxis.set_major_locator(ticker.FixedLocator(positions))
    ax.yaxis.set_major_formatter(ticker.FixedFormatter(labels))
    ax.set_ylim(0, 1.15)
    plt.legend(ncol=3)
    # ax.xaxis.grid(True, linestyle='-.')
    ax.set_title('Happen rate of different results of lazyselect with\ndifferent sampling ratios')
    ax.set_xlabel('sampling ratio')
    positions = sampling_ratios
    labels = ['1/8', '2/8', '3/8', '4/8', '5/8', '6/8', '7/8', '1']
    ax.xaxis.set_major_locator(ticker.FixedLocator(positions))
    ax.xaxis.set_major_formatter(ticker.FixedFormatter(labels))
    # ax.get_xaxis().set_major_formatter(ticker.ScalarFormatter())

    # fig.tight_layout()
    # plt.subplots_adjust(right=0.88)
    # plt.show()
    plt.savefig(_filepath)
    plt.close()

    fig, ax = plt.subplots()
    ax.plot(sampling_ratios, avg_time_consumes, 'o-', linewidth=2.0, color='red')
    ax.set_ylabel('consumption (ms)')
    # ax.tick_params(axis='y', labelcolor='red')
    positions = sampling_ratios
    labels = ['1/8', '2/8', '3/8', '4/8', '5/8', '6/8', '7/8', '1']
    ax.xaxis.set_major_locator(ticker.FixedLocator(positions))
    ax.xaxis.set_major_formatter(ticker.FixedFormatter(labels))
    ax.set_title('Time consumption of lazyselect with\ndifferent sampling ratios')
    ax.set_xlabel('sampling ratio')
    ax.set(ylim=(0, 25), yticks=range(0, 25, 2),)
    # fig.tight_layout()
    plt.grid(True, linestyle='-.')
    # plt.subplots_adjust(right=0.88)
    # plt.show()
    plt.savefig(_filepath_time)
    plt.close()


def save_compare_plot(_results, _labels, _filepath):
    fig, ax = plt.subplots()
    for idx, result in enumerate(_results):
        dataset_sizes = [float(each) for each in result['dataset_sizes']]
        avg_time_consumes = [each * 1000 for each in result['avg_time_consume']]
        ax.plot(dataset_sizes, avg_time_consumes, 'o-', label=_labels[idx], linewidth=2.0)
    ax.set_title('Time consumption of k-th select algorithms with\ndifferent dataset sizes')
    ax.set_xlabel('dataset size')
    ax.set_ylabel('consumption (ms)')
    ax.set(ylim=(0, None))
    ax.legend()
    ax.grid(True, linestyle='-.')
    ax.set_xscale('log')

    ax.set_xticks(_results[0]['dataset_sizes'])
    ax.get_xaxis().set_major_formatter(ticker.ScalarFormatter())

    plt.savefig(_filepath)


if __name__ == '__main__':
    # result_names = ['quickselect', 'lazyselect']
    # labels = ['QuickSelect', 'LazySelect']
    # results = [load_variable(os.path.join('results', each + '.pkl')) for each in result_names]
    # save_compare_plot(results, labels, 'results/result.svg')

    result_name = 'lazyselect_analysis_with_fail'
    result = load_variable(os.path.join('results', result_name + '.pkl'))
    print(result)
    save_lazyselect_analysis_plot(result, 'results/lazyselect_analysis_statistic.png',
                                  'results/lazyselect_analysis_time.png')
